#coding:utf-8
# -*- coding: utf-8 -*-
#@Time : 2020/11/6 11:03
#@Author: 王鹏-wangpeng99
#@File : MyTreeTest.py

import random
from treelib import Node, Tree
import networkx as nx
import matplotlib.pyplot as plt
import matplotlib.colors as colors
import math



tree = Tree()

tmplst=['root']
#写入root节点
tree.create_node('root', 'root')
for i in range(20):
    node = 'node{}'.format(i)
    tree.create_node(node, node, parent=random.choice(tmplst))
    tmplst.append(node)
tree.show()

# tree.create_node("Harry", "harry")  # root node
# tree.create_node("Jane", "jane", parent="harry")
# tree.create_node("Bill", "bill", parent="harry")
# tree.create_node("Diane", "diane", parent="jane")
# tree.create_node("Mary", "mary", parent="diane")
# tree.create_node("Mark", "mark", parent="jane")
# tree.show()

print('='*50)


#节点信息
nodelist= tree.all_nodes()
nodeidentifiers=[item.identifier for item in nodelist]
print(nodelist)
#递归获取所有的节点关联关系
def parseNodeLink(node=None,nodelinklst=None):
    children=tree.children(node)
    if children and len(children)>0:
        for child in children:
            nodelinklst.append((node, child.identifier))
            parseNodeLink(child.identifier, nodelinklst)
    else:
        return
nodelinklst=[]
parseNodeLink(tree.root,nodelinklst)
print(nodelinklst)
print('='*50)




def draw1():
    # 开始绘图
    g = nx.Graph()
    # 增加节点信息
    g.add_nodes_from(
        [(item.identifier, {'tag': item.tag, 'identifier': item.identifier, 'data': item.data}) for item in nodelist])
    # 增加边信息
    g.add_edges_from(nodelinklst)
    # pos信息
    pos = nx.spectral_layout(g)
    # 画图开始
    # nx.draw(g, pos=nx.spectral_layout(g), node_color='r', edge_color='b',with_labels=True)
    # 画root节点
    nx.draw_networkx_nodes(g, pos=pos, nodelist=[tree.root], node_shape='*', node_color='r')
    # 画普通节点
    nx.draw_networkx_nodes(g, pos=pos, nodelist=[item.identifier for item in nodelist[1:]], node_shape='.',
                           node_color=[plt.cm.tab10(idx%10) for idx in range(len(nodelist[1:]))], node_size=1000)
    # 画边信息
    nx.draw_networkx_edges(g, pos=pos, edgelist=nodelinklst, edge_color='b')
    # 画节点标签
    nx.draw_networkx_labels(g, pos=pos, labels={item.identifier: item.identifier for item in nodelist}, font_color='k',
                            font_family='sans-serif')
    plt.show()

    for key, val in pos.items():
        print("key=" + key)
        print("val=" + str(val))
def draw2():
    # 开始绘图
    g = nx.Graph()
    # 增加节点信息
    g.add_nodes_from(
        [(item.identifier, {'tag': item.tag, 'identifier': item.identifier, 'data': item.data}) for item in nodelist])
    # 增加边信息
    g.add_edges_from(nodelinklst)
    # pos信息
    # pos = nx.fruchterman_reingold_layout(g,k=2/math.sqrt(len(nodelist)),iterations=100,)
    pos = nx.fruchterman_reingold_layout(g,iterations=100)
    # 画图开始
    nx.draw(g, pos=pos,
            #节点信息
            nodelist=nodeidentifiers,node_shape='.', node_color=[plt.cm.tab10(idx%10) for idx in range(len(nodelist))], node_size=1000,
            #边信息
            edgelist=nodelinklst, edge_color='b',
            #label信息
            labels={item.identifier: item.identifier for item in nodelist}, font_color='k',font_family='sans-serif',font_size='8')
    plt.show()
if __name__ == '__main__':
    # draw1()
    draw2()